import streamlit as st
import pandas as pd
import numpy as np
from datetime import datetime
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import matplotlib.pyplot as plt
import seaborn as sns
from io import BytesIO
import base64

# 设置页面配置
st.set_page_config(
    page_title="新冠疫情数据分析",
    page_icon="🦠",
    layout="wide"
)


# 数据预处理函数
def preprocess_data(df):
    # 转换日期格式
    df['date'] = pd.to_datetime(df['dateId'].astype(str), format='%Y%m%d')
    # 按日期排序
    df = df.sort_values('date')
    # 处理空值
    df = df.fillna(0)
    # 处理异常值（将负值设为0）
    numeric_columns = df.select_dtypes(include=[np.number]).columns
    for col in numeric_columns:
        df[col] = df[col].clip(lower=0)
    return df


# 数据验证函数
def validate_data(df):
    required_columns = [
        'dateId', 'confirmedCount', 'confirmedIncr',
        'curedCount', 'curedIncr', 'currentConfirmedCount',
        'currentConfirmedIncr', 'deadCount', 'deadIncr'
    ]
    return all(col in df.columns for col in required_columns)


# 下载图表函数
def get_download_link(fig, filename):
    buf = BytesIO()
    fig.write_image(buf, format="png")
    buf.seek(0)
    b64 = base64.b64encode(buf.read()).decode()
    href = f'<a href="data:image/png;base64,{b64}" download="{filename}">下载图表</a>'
    return href


def main():
    st.title("🦠 新冠疫情数据分析平台")
    st.markdown("---")

    # 文件上传部分
    st.sidebar.header("数据上传")
    uploaded_file = st.sidebar.file_uploader(
        "请上传CSV格式的疫情数据文件",
        type=['csv'],
        help="支持的文件格式：china_provincedata.csv, countrydata.csv, world_total_data.csv"
    )

    if uploaded_file is not None:
        try:
            df = pd.read_csv(uploaded_file)
            if not validate_data(df):
                st.error("数据格式不正确！请确保包含所有必需的列。")
                return

            # 数据预处理
            df = preprocess_data(df)

            # 侧边栏控件
            st.sidebar.header("分析参数设置")

            # 时间范围选择
            date_range = st.sidebar.date_input(
                "选择时间范围",
                [df['date'].min(), df['date'].max()],
                min_value=df['date'].min().to_pydatetime(),
                max_value=df['date'].max().to_pydatetime()
            )

            # 感染阈值设置
            threshold = st.sidebar.slider(
                "设置感染人数阈值",
                min_value=0,
                max_value=int(df['confirmedCount'].max()),
                value=int(df['confirmedCount'].max() * 0.7)
            )

            # 数据分析和可视化
            st.header("数据分析结果")

            # 1. 总体趋势分析
            col1, col2 = st.columns(2)

            with col1:
                st.subheader("累计确诊病例趋势")
                fig_total = px.line(
                    df,
                    x='date',
                    y='confirmedCount',
                    title='累计确诊病例随时间变化趋势'
                )
                st.plotly_chart(fig_total, use_container_width=True)

            with col2:
                st.subheader("新增确诊病例趋势")
                fig_new = px.bar(
                    df,
                    x='date',
                    y='confirmedIncr',
                    title='每日新增确诊病例趋势'
                )
                st.plotly_chart(fig_new, use_container_width=True)

            # 2. 高风险区域分析
            st.subheader("高风险区域分析")
            if 'provinceName' in df.columns:  # 针对省份数据
                high_risk = df[df['confirmedCount'] > threshold]['provinceName'].unique()
            elif 'countryName' in df.columns:  # 针对国家数据
                high_risk = df[df['confirmedCount'] > threshold]['countryName'].unique()

            st.warning(f"当前有 {len(high_risk)} 个地区超过阈值，被标记为高风险区域")
            st.write("高风险区域列表：", high_risk)

            # 3. 多维度分析
            st.subheader("多维度分析")
            fig_multi = make_subplots(
                rows=2, cols=2,
                subplot_titles=(
                    "确诊/治愈/死亡累计数据对比",
                    "新增确诊病例分布",
                    "治愈率和死亡率变化",
                    "现存确诊病例变化"
                )
            )

            # 添加确诊/治愈/死亡累计数据
            fig_multi.add_trace(
                go.Scatter(x=df['date'], y=df['confirmedCount'], name="确诊"),
                row=1, col=1
            )
            fig_multi.add_trace(
                go.Scatter(x=df['date'], y=df['curedCount'], name="治愈"),
                row=1, col=1
            )
            fig_multi.add_trace(
                go.Scatter(x=df['date'], y=df['deadCount'], name="死亡"),
                row=1, col=1
            )

            # 添加新增确诊病例分布
            fig_multi.add_trace(
                go.Histogram(x=df['confirmedIncr'], name="新增确诊分布"),
                row=1, col=2
            )

            # 添加治愈率和死亡率
            df['cure_rate'] = df['curedCount'] / df['confirmedCount'] * 100
            df['death_rate'] = df['deadCount'] / df['confirmedCount'] * 100

            fig_multi.add_trace(
                go.Scatter(x=df['date'], y=df['cure_rate'], name="治愈率"),
                row=2, col=1
            )
            fig_multi.add_trace(
                go.Scatter(x=df['date'], y=df['death_rate'], name="死亡率"),
                row=2, col=1
            )

            # 添加现存确诊病例
            fig_multi.add_trace(
                go.Scatter(x=df['date'], y=df['currentConfirmedCount'], name="现存确诊"),
                row=2, col=2
            )

            fig_multi.update_layout(height=800, showlegend=True)
            st.plotly_chart(fig_multi, use_container_width=True)

            # 下载按钮
            st.markdown("---")
            st.subheader("导出分析结果")
            if st.button("生成分析报告"):
                # 创建分析报告
                report = f"""
                # 新冠疫情数据分析报告

                ## 分析时间范围
                - 开始日期：{date_range[0]}
                - 结束日期：{date_range[1]}

                ## 高风险区域（超过阈值 {threshold}）
                {', '.join(high_risk)}

                ## 关键指标
                - 最大单日新增：{df['confirmedIncr'].max()}
                - 累计确诊总数：{df['confirmedCount'].max()}
                - 当前治愈率：{df['cure_rate'].iloc[-1]:.2f}%
                - 当前死亡率：{df['death_rate'].iloc[-1]:.2f}%
                """

                # 创建下载链接
                st.download_button(
                    label="下载分析报告",
                    data=report,
                    file_name="covid19_analysis_report.txt",
                    mime="text/plain"
                )

        except Exception as e:
            st.error(f"处理数据时出错：{str(e)}")

    else:
        st.info("请上传数据文件开始分析")


if __name__ == "__main__":
    main()